import vtk
import numpy as np


def create_sample_structured_grid():
    """创建示例结构化网格数据"""
    # 定义网格尺寸
    grid = vtk.vtkStructuredGrid()
    grid.SetDimensions(20, 20, 20)

    # 创建点坐标
    points = vtk.vtkPoints()

    # 创建速度矢量数据
    vectors = vtk.vtkFloatArray()
    vectors.SetNumberOfComponents(3)
    vectors.SetName("Velocity")

    points_list = []
    velocities = []
    # 生成网格数据和速度场

    for k in range(20):
        for j in range(20):
            for i in range(20):
                x = i / 19.0 * 4.0 - 2.0
                y = j / 19.0 * 4.0 - 2.0
                z = k / 19.0 * 4.0 - 2.0

                points_list.append([x, y, z])


                # 创建涡旋速度场
                vx = -y
                vy = x
                vz = 0.1 * z

                velocities.append([vx, vy, vz])

                if (x ** 2 + y ** 2 + z ** 2) <= 12:
                    points.InsertNextPoint(x, y, z)
                    vectors.InsertNextTuple3(vx, vy, vz)

    points_numpy_array = np.array(points_list)
    points_ = vtk.vtkPoints()
    points_.SetNumberOfPoints(len(points_numpy_array))

    # vtk_point_data = vtk.vtkFloatArray()
    # vtk_point_data.SetVoidArray(points_numpy_array.flatten())
    # points_.SetData(vtk_point_data)

    grid.SetPoints(points)
    grid.GetPointData().SetVectors(vectors)

    return grid


def main():
    # 示例数据
    names = np.array(['Apple', 'Orange', 'Pear', 'Apple', 'Pear', 'Orange'])
    sales = np.array([5, 4, 8, 7, 2, 1])

    # 合并数组并按水果名称排序
    group = np.array(list(zip(names, sales)))
    group = group[np.argsort(group[:, 0])]

    # 创建示例结构化网格
    structured_grid = create_sample_structured_grid()

    # 创建流线种子点
    seeds = vtk.vtkPointSource()
    seeds.SetCenter(0, 0, 0)
    seeds.SetRadius(1.5)
    seeds.SetNumberOfPoints(50)

    # 创建流线追踪器
    stream_tracer = vtk.vtkStreamTracer()
    stream_tracer.SetInputData(structured_grid)
    stream_tracer.SetSourceConnection(seeds.GetOutputPort())
    stream_tracer.SetIntegrationDirectionToBoth()  # 双向积分
    stream_tracer.SetMaximumPropagation(100)  # 最大传播长度
    stream_tracer.SetIntegrationStepUnit(2)  # 单元长度
    stream_tracer.SetInitialIntegrationStep(0.2)
    stream_tracer.SetMaximumNumberOfSteps(2000)

    # 使用高阶积分方法提高精度
    integrator = vtk.vtkRungeKutta45()
    stream_tracer.SetIntegrator(integrator)

    # 创建流线映射器
    stream_mapper = vtk.vtkPolyDataMapper()
    stream_mapper.SetInputConnection(stream_tracer.GetOutputPort())
    stream_mapper.SetScalarRange(structured_grid.GetScalarRange())

    # 创建流线演员
    stream_actor = vtk.vtkActor()
    stream_actor.SetMapper(stream_mapper)
    stream_actor.GetProperty().SetLineWidth(2.0)

    # 创建网格轮廓
    outline = vtk.vtkStructuredGridOutlineFilter()
    outline.SetInputData(structured_grid)

    outline_mapper = vtk.vtkPolyDataMapper()
    outline_mapper.SetInputConnection(outline.GetOutputPort())

    outline_actor = vtk.vtkActor()
    outline_actor.SetMapper(outline_mapper)
    outline_actor.GetProperty().SetColor(1, 1, 1)  # 白色轮廓

    # 创建渲染器
    renderer = vtk.vtkRenderer()
    renderer.AddActor(stream_actor)
    renderer.AddActor(outline_actor)
    renderer.SetBackground(0.1, 0.1, 0.3)  # 深蓝色背景

    # 创建渲染窗口
    render_window = vtk.vtkRenderWindow()
    render_window.AddRenderer(renderer)
    render_window.SetSize(800, 600)
    render_window.SetWindowName("VTK Stream Tracer with Structured Grid")

    # 创建交互器
    interactor = vtk.vtkRenderWindowInteractor()
    interactor.SetRenderWindow(render_window)

    # 设置相机位置
    camera = renderer.GetActiveCamera()
    camera.SetPosition(10, 10, 10)
    camera.SetFocalPoint(0, 0, 0)

    # 开始渲染和交互
    render_window.Render()
    interactor.Start()


if __name__ == "__main__":
    main()
